package com.ifp.opengate.boot.factory;

import java.time.Instant;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.validation.constraints.Min;
import javax.validation.constraints.NotNull;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.BeansException;
import org.springframework.cloud.gateway.event.FilterArgsEvent;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationUtils;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.validation.Validator;
import org.springframework.validation.annotation.Validated;

import com.ifp.core.log.Trace;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * 项目名：ifp-gateway <br>
 * 包名：com.ifp.opengate.boot.factory <br>
 * 类名：IfpRedisRateLimiter <br>
 * 描述：
 *
 * @author Suremotoo
 * @create 2021-06-24 16:31
 */

public class IfpRedisRateLimiter extends AbstractRateLimiter<IfpRedisRateLimiter.Config>
    implements ApplicationContextAware {

    public static final String REPLENISH_RATE_KEY = "replenishRate";

    public static final String BURST_CAPACITY_KEY = "burstCapacity";
    public static final String CONFIGURATION_PROPERTY_NAME = "redis-rate-limiter";
    public static final String REDIS_SCRIPT_NAME = "redisRequestRateLimiterScript";
    public static final String REMAINING_HEADER = "X-RateLimit-Remaining";
    public static final String REPLENISH_RATE_HEADER = "X-RateLimit-Replenish-Rate";
    public static final String BURST_CAPACITY_HEADER = "X-RateLimit-Burst-Capacity";
    private Log log = LogFactory.getLog(this.getClass());
    private ReactiveRedisTemplate<String, String> redisTemplate;
    private RedisScript<List<Long>> script;
    private AtomicBoolean initialized = new AtomicBoolean(false);
    private IfpRedisRateLimiter.Config defaultConfig;
    private boolean includeHeaders = true;
    private String remainingHeader = "X-RateLimit-Remaining";
    private String replenishRateHeader = "X-RateLimit-Replenish-Rate";
    private String burstCapacityHeader = "X-RateLimit-Burst-Capacity";

    public IfpRedisRateLimiter(ReactiveRedisTemplate<String, String> redisTemplate, RedisScript<List<Long>> script,
        Validator validator) {
        super(IfpRedisRateLimiter.Config.class, "redis-rate-limiter", validator);
        this.redisTemplate = redisTemplate;
        this.script = script;
        this.initialized.compareAndSet(false, true);
    }

    public IfpRedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity) {
        super(IfpRedisRateLimiter.Config.class, "redis-rate-limiter", (Validator)null);
        this.defaultConfig = (new IfpRedisRateLimiter.Config()).setReplenishRate(defaultReplenishRate)
            .setBurstCapacity(defaultBurstCapacity);
    }

    static List<String> getKeys(String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    public boolean isIncludeHeaders() {
        return this.includeHeaders;
    }

    public void setIncludeHeaders(boolean includeHeaders) {
        this.includeHeaders = includeHeaders;
    }

    public String getRemainingHeader() {
        return this.remainingHeader;
    }

    public void setRemainingHeader(String remainingHeader) {
        this.remainingHeader = remainingHeader;
    }

    public String getReplenishRateHeader() {
        return this.replenishRateHeader;
    }

    public void setReplenishRateHeader(String replenishRateHeader) {
        this.replenishRateHeader = replenishRateHeader;
    }

    public String getBurstCapacityHeader() {
        return this.burstCapacityHeader;
    }

    public void setBurstCapacityHeader(String burstCapacityHeader) {
        this.burstCapacityHeader = burstCapacityHeader;
    }

    @Override
    public void setApplicationContext(ApplicationContext context) throws BeansException {
        if (this.initialized.compareAndSet(false, true)) {
            this.redisTemplate =
                (ReactiveRedisTemplate)context.getBean("stringReactiveRedisTemplate", ReactiveRedisTemplate.class);
            this.script = (RedisScript)context.getBean("redisRequestRateLimiterScript", RedisScript.class);
            if (context.getBeanNamesForType(Validator.class).length > 0) {
                this.setValidator((Validator)context.getBean(Validator.class));
            }
        }

    }

    IfpRedisRateLimiter.Config getDefaultConfig() {
        return this.defaultConfig;
    }

    @Override
    public void onApplicationEvent(FilterArgsEvent event) {
        Map<String, Object> args = event.getArgs();
        if (!args.isEmpty() && this.hasRelevantKey(args)) {
            String routeId = event.getRouteId();
            Config routeConfig = this.newConfig();
            ConfigurationUtils.bind(routeConfig, args, super.getConfigurationPropertyName(),
                this.getConfigurationPropertyName(), this.getValidator());
            // 这里重写 id，防止过冲规则重复被覆盖 by liulx
            // 使用routeId+KeyResolver.hashcode作为配置id，防止重复
            routeId = routeId + event.getArgs().get("key-resolver").hashCode();

            this.getConfig().put(routeId, routeConfig);
        }
    }

    private boolean hasRelevantKey(Map<String, Object> args) {
        return args.keySet().stream().anyMatch((key) -> {
            return key.startsWith(this.getConfigurationPropertyName() + ".");
        });
    }

    @Override
    public Mono<Response> isAllowed(String routeId, String id) {
        if (!this.initialized.get()) {
            throw new IllegalStateException("IfpRedisRateLimiter is not initialized");
        } else {
            IfpRedisRateLimiter.Config routeConfig = this.loadConfiguration(routeId);
            int replenishRate = routeConfig.getReplenishRate();
            int burstCapacity = routeConfig.getBurstCapacity();

            try {
                String tempId = UUID.randomUUID().toString();
                Trace.logInfo(Trace.MODULE_ACTION, "IfpRedisRateLimiter {}  start time {}", tempId,
                    System.currentTimeMillis());
                List<String> keys = getKeys(id);
                List<String> scriptArgs =
                    Arrays.asList(replenishRate + "", burstCapacity + "", Instant.now().getEpochSecond() + "", "1");
                Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
                Trace.logInfo(Trace.MODULE_ACTION, "IfpRedisRateLimiter {} execute time {}", tempId,
                    System.currentTimeMillis());
                return flux.onErrorResume((throwable) -> {
                    return Flux.just(Arrays.asList(1L, -1L));
                }).reduce(new ArrayList(), (longs, l) -> {
                    longs.addAll(l);
                    return longs;
                }).map((results) -> {
                    boolean allowed = (Long)results.get(0) == 1L;
                    Long tokensLeft = (Long)results.get(1);
                    Response response = new Response(allowed, this.getHeaders(routeConfig, tokensLeft));
                    if (this.log.isDebugEnabled()) {
                        this.log.debug("response: " + response);
                    }
                    Trace.logInfo(Trace.MODULE_ACTION, "IfpRedisRateLimiter {} execute after time {}", tempId,
                        System.currentTimeMillis());
                    return response;
                });
            } catch (Exception var9) {
                this.log.error("Error determining if user allowed from redis", var9);
                return Mono.just(new Response(true, this.getHeaders(routeConfig, -1L)));
            }
        }
    }

    IfpRedisRateLimiter.Config loadConfiguration(String routeId) {
        IfpRedisRateLimiter.Config routeConfig =
            (IfpRedisRateLimiter.Config)this.getConfig().getOrDefault(routeId, this.defaultConfig);
        if (routeConfig == null) {
            routeConfig = (IfpRedisRateLimiter.Config)this.getConfig().get("defaultFilters");
        }

        if (routeConfig == null) {
            throw new IllegalArgumentException("No Configuration found for route " + routeId + " or defaultFilters");
        } else {
            return routeConfig;
        }
    }

    @NotNull
    public Map<String, String> getHeaders(IfpRedisRateLimiter.Config config, Long tokensLeft) {
        Map<String, String> headers = new HashMap();
        if (this.isIncludeHeaders()) {
            headers.put(this.remainingHeader, tokensLeft.toString());
            headers.put(this.replenishRateHeader, String.valueOf(config.getReplenishRate()));
            headers.put(this.burstCapacityHeader, String.valueOf(config.getBurstCapacity()));
        }

        return headers;
    }

    @Validated
    public static class Config {
        @Min(1L)
        private int replenishRate;
        @Min(1L)
        private int burstCapacity = 1;

        public Config() {}

        public int getReplenishRate() {
            return this.replenishRate;
        }

        public IfpRedisRateLimiter.Config setReplenishRate(int replenishRate) {
            this.replenishRate = replenishRate;
            return this;
        }

        public int getBurstCapacity() {
            return this.burstCapacity;
        }

        public IfpRedisRateLimiter.Config setBurstCapacity(int burstCapacity) {
            this.burstCapacity = burstCapacity;
            return this;
        }

        @Override
        public String toString() {
            return "Config{replenishRate=" + this.replenishRate + ", burstCapacity=" + this.burstCapacity + '}';
        }
    }
}
